# RAG流程 - 环节6：生成回答
# 功能：基于检索到的相关片段，调用大模型生成准确、可溯源的回答
# 版本：适配通义千问大模型，支持保险文档专业问答场景

import time
from typing import List, Dict
import dashscope
from dashscope import Generation
from hybrid_retrieval import HybridRetriever  # 导入环节5的混合检索器

# 配置参数（请根据实际环境修改）
CONFIG = {
    # 通义千问API配置
    "tongyi_api_key": "sk-6434b554122c4b049ceb805d703f695b",  # 你的阿里云API Key
    "llm_model": "qwen-plus",  # 大模型选择：qwen-plus/qwen-max
    # 生成参数（控制回答风格）
    "max_tokens": 1024,  # 最大回答长度
    "temperature": 0.2,  # 随机性（0.1-0.3适合专业问答）
    "top_p": 0.8,  # 采样阈值
    # 检索配置（与环节5保持一致）
    "es_host": "localhost",
    "es_port": 9200,
    "es_index": "insurance_manual",
    "milvus_host": "localhost",
    "milvus_port": "19530",
    "milvus_collection": "insurance_manual",
    "top_k": 3  # 用于生成回答的相关片段数量
}


class AnswerGenerator:
    def __init__(self):
        """初始化回答生成器，包含检索器和大模型组件"""
        # 初始化混合检索器（环节5）
        self.retriever = HybridRetriever()

        # 初始化大模型
        self._init_llm()

    def _init_llm(self) -> None:
        """初始化通义千问大模型"""
        try:
            dashscope.api_key = CONFIG["tongyi_api_key"]

            # 验证API有效性
            test_response = Generation.call(
                model=CONFIG["llm_model"],
                prompt="验证连接",
                max_tokens=10
            )

            if test_response.status_code != 200:
                raise Exception(f"API验证失败: {test_response.message}")

            print(f"✅ 大模型初始化成功 (模型: {CONFIG['llm_model']})")
        except Exception as e:
            raise Exception(f"大模型初始化失败: {str(e)}\n请检查API Key和网络连接")

    def build_prompt(self, query: str, relevant_docs: List[Dict]) -> str:
        """
        构建提示词，指导大模型生成符合要求的回答
        参数:
            query: 用户问题
            relevant_docs: 检索到的相关文档片段
        返回: 结构化提示词
        """
        # 格式化检索到的文档片段
        formatted_docs = ""
        for i, doc in enumerate(relevant_docs[:CONFIG["top_k"]], 1):
            formatted_docs += f"""
【相关片段 {i}】
来源页码: 第{doc['page']}页
内容: {doc['text'][:500]}  # 截断过长内容，避免提示词溢出
"""

        # 提示词模板（核心规则定义）
        prompt = f"""
你是专业的保险文档问答助手，需严格按照以下规则回答用户问题:

1. 仅使用【相关片段】中的信息，绝对不能编造未提及的内容
2. 回答需分点清晰，使用数字编号（1. 2. 3.）
3. 每个要点必须标注来源页码（格式: [来源: 第X页]）
4. 若多个片段包含相同信息，合并说明并标注所有来源
5. 若没有相关信息，直接回复"抱歉，未找到相关信息"

【相关片段】:
{formatted_docs}

【用户问题】: {query}

【回答】:
"""
        return prompt.strip()

    def generate_answer(self, prompt: str) -> str:
        """调用大模型生成回答"""
        try:
            start_time = time.time()
            print(f"📤 正在生成回答...")

            response = Generation.call(
                model=CONFIG["llm_model"],
                prompt=prompt,
                max_tokens=CONFIG["max_tokens"],
                temperature=CONFIG["temperature"],
                top_p=CONFIG["top_p"],
                result_format="text"
            )

            if response.status_code == 200:
                cost_time = time.time() - start_time
                print(f"✅ 回答生成完成 (耗时: {cost_time:.2f}秒)")
                return response.output["text"].strip()
            else:
                raise Exception(f"大模型调用失败: {response.message}")
        except Exception as e:
            raise Exception(f"生成回答时出错: {str(e)}")

    def process_query(self, query: str) -> None:
        """完整处理流程: 检索→构建提示词→生成回答→展示结果"""
        try:
            # 1. 检索相关片段
            print(f"\n🔍 检索: {query}")
            relevant_docs = self.retriever.hybrid_search(query)

            # 2. 构建提示词
            prompt = self.build_prompt(query, relevant_docs)

            # 3. 生成回答
            answer = self.generate_answer(prompt)

            # 4. 展示结果
            self._display_result(query, answer)
        except Exception as e:
            print(f"❌ 处理失败: {str(e)}")

    def _display_result(self, query: str, answer: str) -> None:
        """格式化显示结果"""
        print("\n" + "=" * 80)
        print(f"【用户问题】: {query}")
        print("-" * 80)
        print(f"【回答】:\n{answer}")
        print("=" * 80)


def main():
    """主函数: 初始化组件并处理测试问题"""
    try:
        print("===== 初始化RAG回答生成系统 =====")
        generator = AnswerGenerator()

        # 测试问题列表
        test_queries = [
            "平安福2024版的投保年龄限制是什么？",
            "健康告知需要提供哪些信息？",
            "保险理赔的基本流程是什么？",
            "企业财产险保障哪些类型的损失？",
            "平安福2024的缴费方式有哪些选择？"
        ]

        # 处理每个问题
        for i, query in enumerate(test_queries, 1):
            print(f"\n===== 处理第{i}个问题 =====")
            generator.process_query(query)

    except Exception as e:
        print(f"\n❌ 系统启动失败: {str(e)}")


if __name__ == "__main__":
    main()
